#!/usr/bin/env python3
#-*- coding:utf8 -*-
# Power by 2020-06-13 14:45:23

import os
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import json
import gzip
import random
import numpy as np
from PIL import Image

def load_data(model='train'):
    """TODO: Docstring for load_data.

    :model: TODO
    :returns: TODO

    """
    datafile='./mnist.json.gz' 
    print("loading data from :{}".format(datafile))
    data=json.load(gzip.open(datafile))
    train_set,val_set,eval_data=data
    IMG_ROWS=28
    IMG_COLS=28
    if model =='train':
        imgs=train_set[0]
        labels=train_set[1]
    elif model =='valid':
        imgs=val_set[0]
        labels=val_set[1]
    elif model=='eval':
        imgs=eval_data[0]
        labels=eval_data[1]
    imgs_length=len(imgs)
    assert len(imgs) == len(labels),"length of imgs should be the same as labels".format(len(imgs),len(labels))
    index_list=list(range(imgs_length))
    batchsize=100
    def data_generator():
        if model =='train':
            random.shuffle(index_list)
        imgs_list=[]
        labels_list=[]
        for i in index_list:
            img=np.reshape(imgs[i],[1,IMG_ROWS,IMG_COLS]).astype('float32')
            label=np.reshape(labels[i],[1]).astype('int64')

            imgs_list.append(img)
            labels_list.append(label)
            if len(imgs_list) == batchsize:
                yield np.array(imgs_list),np.array(labels_list)
                imgs_list=[]
                labels_list=[]
        if len(imgs_list) >0:
            yield np.array(imgs_list),np.array(labels_list)
    return data_generator

